# __author__ = 'heyin'
# __date__ = '2018/12/26 15:43'
from pyecharts import Line, Scatter
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_wine
from sklearn.tree import export_graphviz


def jueceshu():
    wine = load_wine()
    # 数据集的详细内容
    # print(wine.feature_names)
    # print(wine.target)
    # print(wine.data)

    x_train, x_test, y_train, y_test = train_test_split(wine.data, wine.target, test_size=0.25, random_state=1)
    # for i in ['gini', 'entropy']:
    #     dtc = DecisionTreeClassifier(criterion=i)
    scores = list()
    for i in range(1, 11):
        print(i)
        dtc = DecisionTreeClassifier(criterion='gini', max_depth=i, random_state=1)
        dtc.fit(x_train, y_train)
        score = dtc.score(x_test, y_test)
        scores.append(score)
        # print(score, '训练集效果：%s' % dtc.score(x_train, y_train))
    l = Scatter()
    l.add('得分与max_depth的关系图', list(range(1, 11)), scores)
    l.render('树.html')



    # 绘图
    # feature_names = ['酒精', '苹果酸', '灰', '灰的碱性', '镁', '总酚', '类黄酮', '非黄烷类酚类', '花青素', '颜色强度', '色调', 'od280 / od315稀释葡萄酒',
    #                 '脯氨酸']
    # feature_names = wine.feature_names
    # class_names = ["one", "tow", "three"]
    # export_graphviz(dtc, out_file='./wine.dot', feature_names=feature_names, class_names=class_names, filled=True)

    # print(*zip(feature_names, dtc.feature_importances_))


if __name__ == '__main__':
    jueceshu()
